Skip to content

Harden model validation: reject __setstate__ hooks and forbid inductor::_reinterpret_tensor#3078

Open
edsavage wants to merge 12 commits into
elastic:mainfrom
edsavage:security/setstate-load-timing
Open

Harden model validation: reject __setstate__ hooks and forbid inductor::_reinterpret_tensor#3078
edsavage wants to merge 12 commits into
elastic:mainfrom
edsavage:security/setstate-load-timing

Conversation

@edsavage

@edsavage edsavage commented Jul 20, 2026

Copy link
Copy Markdown
Contributor

Summary

Two related TorchScript model-validation hardening fixes, landed together as one coherent change (supersedes #3079). Both address privately reported security findings; details are tracked internally.

1. __setstate__ load-time execution gap

  • torch::jit::load runs a module's __setstate__ during deserialization, before the post-load graph validator (which only inspects the forward graph) could ever run. A forbidden op hidden in __setstate__ therefore executes at load and is invisible to the validator.
  • Remediation (matches the recommended one): statically scan the archive and refuse any model that embeds __setstate__/__getstate__ before load, without executing model code.

2. inductor::_reinterpret_tensor heap-OOB bypass

  • After aten::as_strided was forbidden, torch.ops.inductor._reinterpret_tensor gives the same class of unchecked storage-offset OOB heap read/write.
  • Add it to FORBIDDEN_OPERATIONS so the post-load graph validator fail-fasts (authoritative, syntax-independent). An op hidden in __setstate__ is already refused pre-load by the hook ban above.

Post-load allowlist / forbid validation is otherwise unchanged.

Test plan

  • Unit tests: custom state hooks rejected; forward-only malicious models have no hooks (caught post-load); inductor::_reinterpret_tensor rejected post-load; benign e5_with_norm.pt accepted; garbage input does not crash
  • CI pytorch_inference / model validator suites green
  • Confirm production-style scripted transformers do not embed custom state hooks

Model-zoo compatibility evidence (blanket hook ban)

Models exported exactly as ES ingests them (eland 9.2.0 TransformerModel.save()torch.jit.trace/save), each archive scanned with the same logic as scanArchiveForCustomStateHooks. Scanner validated by a positive control (explicit __setstate__ correctly rejected).

Family Model Task(s) quant=False quant=True
BERT bert-base-uncased (full) text_embedding, fill_mask accepted accepted
DistilBERT distilbert-base-uncased text_embedding accepted accepted
RoBERTa all-distilroberta-v1 text_embedding accepted accepted
ELECTRA electra-small-discriminator text_embedding accepted accepted
DeBERTa-v2 deberta-v3-xsmall fill_mask accepted accepted
MPNet all-mpnet-base-v2 text_embedding accepted accepted

Plus all local fixtures (real E5 embedding, norm model, the three ES integration-test models). The --quantize path (the main theoretical risk) is clean: TorchScript serialises quantized packed params via a bound C++ mechanism, so the hook literals never appear in the archive.

Related: #3080 (seccomp arch/ABI check), #3081 (controller non-dumpable).

TorchScript runs a module's __setstate__ during torch::jit::load(), before the
loaded module reaches CModelGraphValidator (which only walks the inlined forward
graph). A forbidden op hidden in __setstate__ is therefore invisible to the
validator and has already executed by the time validation would run.

Add fixtures and a self-contained repro that demonstrate the gap:
- generate_malicious_models.py: add SetStateFileReaderModel and a submodule
  variant (aten::from_file in __setstate__, benign forward); surface
  __setstate__ ops in the generator output.
- test/test_setstate_load_timing.py: torch-only, build-free repro proving
  load-time execution (benign print probe) and op-hiding (static inspection).
- test_pytorch_inference_evil_models.py: register the fixtures as regression
  guards for the forthcoming pre-load scan.

No fix yet; this only reproduces and documents the vulnerability.

Co-authored-by: Cursor <[email protected]>
Close the load-time execution gap: torch::jit::load runs a module's
__setstate__ during deserialization, before CModelGraphValidator (which walks
an already-loaded module's graph) can run. A forbidden op hidden in
__setstate__ therefore executes at load time and never appears in the forward
graph the validator inspects.

Add CModelGraphValidator::scanSerialisedCodeForForbiddenOps, which uses a
PyTorchStreamReader over the buffered archive bytes to textually scan the
serialised TorchScript code (code/**/*.py, including every submodule's
__setstate__) for the forbidden aten ops, WITHOUT loading the model. Main.cc
runs this scan before torch::jit::load and rejects via HANDLE_FATAL, so no
model code executes. Expose the buffered bytes via CBufferedIStreamAdapter::buffer().

This is a defense-in-depth textual check targeted at the curated forbidden set
(from_file, as_strided, save); the post-load allowlist validation is unchanged,
and seccomp remains the syscall backstop. Unit tests cover both setstate
fixtures, forward-graph attacks, a benign model (no false positive), and
malformed input; the binary integration test now expects the setstate models to
be rejected with the forbidden-operations message.

Co-authored-by: Cursor <[email protected]>
@edsavage edsavage changed the title Reproduce __setstate__ load-time execution gap in model validation Reproduce and fix __setstate__ load-time execution gap in model validation Jul 20, 2026
@elasticsearchmachine

Copy link
Copy Markdown

Hi @edsavage, I've created a changelog YAML for you.

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

This PR closes a security-relevant validation gap in bin/pytorch_inference: TorchScript can execute a module’s __setstate__ during torch::jit::load(), so forbidden ops hidden there could run before the existing post-load graph validation. The fix adds a pre-load static scan of the serialized TorchScript source within the .pt archive (including __setstate__) and expands test coverage to reproduce and prevent regressions.

Changes:

  • Add CModelGraphValidator::scanSerialisedCodeForForbiddenOps() to scan code/**/*.py inside a serialized .pt archive for forbidden aten ops without loading the model.
  • Invoke the pre-load scan in pytorch_inference before torch::jit::load() (when model validation is enabled).
  • Add/extend Python + C++ tests and model-generation tooling to reproduce the gap and assert rejection of __setstate__-hidden forbidden ops.

Reviewed changes

Copilot reviewed 4 out of 11 changed files in this pull request and generated 1 comment.

Show a summary per file
File Description
test/test_setstate_load_timing.py New standalone repro script demonstrating __setstate__ executes during load and is invisible to forward-graph-only validation.
test/test_pytorch_inference_evil_models.py Adds __setstate__-based “evil models” and updates expectations to require clean rejection.
docs/changelog/3078.yaml Changelog entry for the security/validation bug fix.
dev-tools/generate_malicious_models.py Adds __setstate__ attack fixtures and surfaces __setstate__ ops during fixture generation.
bin/pytorch_inference/unittest/CModelGraphValidatorTest.cc Adds unit tests for the new pre-load scan (setstate fixtures, benign model, garbage input).
bin/pytorch_inference/Main.cc Runs the pre-load serialized-code scan before calling torch::jit::load().
bin/pytorch_inference/CModelGraphValidator.h Declares the new serialized-code scanning API and documents rationale/behavior.
bin/pytorch_inference/CModelGraphValidator.cc Implements in-memory archive parsing via PyTorchStreamReader and textual scanning of code/**/*.py.
bin/pytorch_inference/CBufferedIStreamAdapter.h Exposes a read-only buffer() view so the already-buffered archive can be scanned pre-load.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Comment thread test/test_setstate_load_timing.py
edsavage and others added 3 commits July 21, 2026 10:21
Address Copilot review: the benign probe must load cleanly for the marker to
prove load-time execution. Require returncode == 0 and that LOAD_RETURNED was
printed, and print returncode + stderr tail on failure, so a crash that happens
to print the marker first is no longer a false positive.

Co-authored-by: Cursor <[email protected]>
Match the #12621 reporter remediation by refusing archives that embed
__setstate__/__getstate__, and keep the forbidden-op pre-scan as defense
in depth. Allowlisted-op RCE gadgets in setstate are otherwise missed.

Co-authored-by: Cursor <[email protected]>
@edsavage edsavage changed the title Reproduce and fix __setstate__ load-time execution gap in model validation Reject TorchScript custom state hooks before load (__setstate__ gap) Jul 21, 2026
@edsavage

Copy link
Copy Markdown
Contributor Author

Updated after comparing with the reporter's 0001-reject-torchscript-state-hooks.patch: a forbidden-op-only pre-scan would miss the real #12621 setstate RCE (allowlisted ops + memory gadgets). We now reject __setstate__/__getstate__ substrings in the archive before load, and retain the forbidden-op scan as defense in depth.

edsavage and others added 2 commits July 21, 2026 15:45
Custom state-hook rejection is enough for the load-time execution gap;
forward-graph ops remain covered by post-load validate().

Co-authored-by: Cursor <[email protected]>
@edsavage

Copy link
Copy Markdown
Contributor Author

Model-zoo compatibility check for the blanket __setstate__/__getstate__ ban

Ran the third test-plan item. Models were exported exactly as ES ingests them (eland 9.2.0 TransformerModel.save()torch.jit.trace/save) and each archive scanned with the same record-substring logic as scanArchiveForCustomStateHooks. Scanner validated by a positive control (a module with an explicit __setstate__ was correctly REJECTED, literal found in code/__torch__.py).

Every legitimate model was accepted (zero false positives):

Family Model Task(s) quant=False quant=True
BERT bert-base-uncased (full) text_embedding, fill_mask accepted accepted
DistilBERT distilbert-base-uncased text_embedding accepted accepted
RoBERTa all-distilroberta-v1 text_embedding accepted accepted
ELECTRA electra-small-discriminator text_embedding accepted accepted
DeBERTa-v2 deberta-v3-xsmall fill_mask accepted accepted
MPNet all-mpnet-base-v2 text_embedding accepted accepted

Plus all local fixtures: real E5 embedding (e5_with_norm.pt), norm_model.pt, and the three ES integration-test models (tiny_text_embedding, tiny_text_expansion, supersimple_pytorch_model_it).

Notable: the --quantize path was the main theoretical risk (dynamic-quantized LinearPackedParams define __getstate__/__setstate__ in Python), but TorchScript serialises packed params via a bound C++ mechanism, so the ASCII literals never appear in the archive — confirmed on a standalone quantized nn.Linear and on all quantized full-model exports above.

Conclusion: the blanket ban does not reject any legitimate model across the encoder families eland/ES supports for NLP.

…, #12242)

Fold the inductor::_reinterpret_tensor hardening (elastic/security#12242) into
the same PR as the __setstate__ load-time hook ban (elastic/security#12621), so
the two model-validation fixes land as one coherent change and supersede elastic#3079.

TorchInductor's _reinterpret_tensor takes an unchecked storage offset and yields
the same class of OOB heap read/write previously mitigated by forbidding
aten::as_strided; it was used to bypass that forbid.

- Add inductor::_reinterpret_tensor to FORBIDDEN_OPERATIONS so the post-load
  graph validator fail-fasts (authoritative, syntax-independent). An op hidden
  in __setstate__ is already refused pre-load by the custom-state-hook ban.
- Add forward-only OOB read/write fixtures + a post-load forbid unit test and
  evil-model integration entries.
- Reconcile the evil-model __setstate__ entries to expect the hook-ban message
  ("custom state hooks") rather than the removed op-scan message.

Co-authored-by: Cursor <[email protected]>
@edsavage edsavage changed the title Reject TorchScript custom state hooks before load (__setstate__ gap) Harden model validation: reject __setstate__ hooks and forbid inductor::_reinterpret_tensor Jul 23, 2026
@edsavage
edsavage marked this pull request as ready for review July 23, 2026 00:48
@elasticsearchmachine

Copy link
Copy Markdown

Pinging @elastic/ml-core (Team:ML)

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Copilot reviewed 4 out of 14 changed files in this pull request and generated no new comments.

@edsavage

Copy link
Copy Markdown
Contributor Author

Test plan item 2 (CI pytorch_inference / model validator suites green) confirmed by build 2844 on head ece0c5d1f — passed on Linux x86_64/aarch64, macOS aarch64 and Windows x86_64.

ml_test_pytorch_inference: 60 tests, 0 failures, 0 errors, including this PR's new cases:

  • testPreLoadScanRejectsCustomStateHooks / ...InSubmodule
  • testPreLoadScanAcceptsBenignModel / ...ForwardOnlyMaliciousModel
  • testPreLoadScanHandlesGarbageInput
  • testMaliciousReinterpretTensorRejectedPostLoad
  • testPrepackedE5ModelWithNorm

All test-plan items are now checked off.

@edsavage

Copy link
Copy Markdown
Contributor Author

buildkite run_pytorch_tests

@edsavage
edsavage requested a review from valeriy42 July 23, 2026 01:37
Comment thread bin/pytorch_inference/CModelGraphValidator.cc Outdated
Comment thread bin/pytorch_inference/unittest/CModelGraphValidatorTest.cc
@edsavage edsavage added auto-backport Automatically merge backport PRs when CI passes v8.19.20 v9.4.5 v9.5.1 labels Jul 23, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants